﻿/*
BF4Rcon.NET, provides a .NET implementation of BF4's RCON interface.
Copyright (C) 2013 Tim Marro

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using NUnit.Framework;
using NUnit.Framework.Constraints;
using TimScripts.BF4Rcon.Frostbite;
using System.Threading;

namespace TimScripts.BF4Rcon.Test
{
    [TestFixture]
    public class RconClientTests
    {
        //wont be able to test the crazy socket stuff
        //just noting since never tried abstracting the socket out this far before

        TestConnection2 Connection;
        RconClient Client;

        [SetUp]
        public void SetUp()
        {
            Connection = new TestConnection2();
            Client = new RconClient(Connection);
        }

        [Test]
        public void ClientOnlyConnectsIfNotConnected()
        {
            Client.Connect();

            Assert.AreEqual(true, Connection.AttemptedConnect);

            Assert.Throws<InvalidOperationException>(() => Client.Connect());
        }

        [Test]
        public void ClientOnlyDisconnectsIfConnected()
        {
            Client.Connect();
            Client.Disconnect();

            Assert.AreEqual(true, Connection.AttemptedDisconnect);

            Connection.ResetFlags();

            Client.Disconnect();

            Assert.AreEqual(false, Connection.AttemptedDisconnect);
        }

        [Test]
        public void CorrectHeaderSet()
        {
            Client.Connect();
            Client.SendRequest("a");

            Assert.AreEqual(true, Connection.LastIsFromClient);
            Assert.AreEqual(false, Connection.LastIsResponse);
        }

        [Test]
        public void SendingWhenNotConnectedThrows()
        {
            Assert.Throws<DisconnectionException>(() => { Client.SendRequest("a"); });
        }

        [Test]
        public void SendingWhenNotConnectedThrows_Async()
        {
            Assert.Throws<DisconnectionException>(async () => { await Client.SendRequestAsync("a"); });
        }

        [Test]
        public void CorrectWordsReturned()
        {
            Connection.EnqueueEvent("wrong");

            Client.Connect();
            var result = Client.SendRequest("a");

            Assert.AreEqual("OK", result[0]);

            Client.Disconnect();
        }


        [Test]
        public async void CorrectWordsReturned_Async()
        {
            Connection.EnqueueEvent("wrong");

            Client.Connect();
            var result = await Client.SendRequestAsync("a");

            Assert.AreEqual("OK", result[0]);

            Client.Disconnect();
        }

        [Test]
        public void WorksWithAlreadyConnectedConnection()
        {
            Connection.Connect();

            RconClient c = new RconClient(Connection);

            var result = c.SendRequest("a");

            Assert.AreEqual("OK", result[0]);
        }

        [Test]
        public async void WorksWithAlreadyConnectedConnection_Async()
        {
            Connection.Connect();

            RconClient c = new RconClient(Connection);

            var result = await c.SendRequestAsync("a");

            Assert.AreEqual("OK", result[0]);
        }

        [Test]
        public void RequestsAccountForExternalConnect()
        {
            Connection.Connect();

            var result = Client.SendRequest("a");

            Assert.AreEqual("OK", result[0]);
        }

        [Test]
        public async void RequestsAccountForExternalConnect_Async()
        {
            Connection.Connect();

            var result = await Client.SendRequestAsync("a");

            Assert.AreEqual("OK", result[0]);
        }

        [Test]
        public void TimeoutThrows()
        {
            Client.Connect();
            Connection.Delay = 2000;

            Assert.Throws<TimeoutException>(() => Client.SendRequest(1, "a"));
        }

        [Test]
        public void CancellationThrows()
        {
            using (CancellationTokenSource cts = new CancellationTokenSource(1000))
            {
                Client.Connect();
                Connection.Delay = 5000;
                Assert.Throws<TaskCanceledException>(async () => await Client.SendRequestAsync(cts.Token, "a"));
            }
        }

        [Test]
        public void ConnectionDisconnectDuringRequestThrows()
        {
            Client.Connect();
            Connection.Delay = 500;

            Task t = Task.Factory.StartNew(() => Client.SendRequest("a"));

            Connection.Disconnect();

            try
            {
                t.Wait();
            }
            catch (AggregateException ex)
            {
                Assert.AreEqual(1, ex.InnerExceptions.Count);
                Assert.That(ex.InnerExceptions[0], Is.TypeOf<DisconnectionException>());
            }
        }

        [Test]
        public void ConnectionDisconnectDuringRequestThrowsAsync()
        {
            Client.Connect();
            Connection.Delay = 500;

            Task t = Client.SendRequestAsync("a");

            Connection.Disconnect();

            try
            {
                t.Wait();
            }
            catch (AggregateException ex)
            {
                Assert.AreEqual(1, ex.InnerExceptions.Count);
                Assert.That(ex.InnerExceptions[0], Is.TypeOf<DisconnectionException>());
            }
        }

        [Test]
        public void SendThrowsWhenConnectionThrows()
        {
            Client.Connect();
            Connection.ThrowSend = true;
            Assert.Throws<TestException>(() => Client.SendRequest("a"));
        }

        [Test]
        public void SendThrowsWhenConnectionThrowsAsync()
        {
            Client.Connect();
            Connection.ThrowSend = true;
            Assert.Throws<TestException>(async () => await Client.SendRequestAsync("a"));
        }

        [Test]
        public void ReceiveThrowsWhenConnectionThrows()
        {
            Client.Connect();
            Connection.Delay = 1000;
            Connection.ThrowReceiveOnce = true;
            Assert.Throws<TestException>(() => Client.SendRequest("a"));
        }

        [Test]
        public void ReceiveThrowsWhenConnectionThrowsAsync()
        {
            IList<string> result;
            Client.Connect();
            Connection.Delay = 2000;
            Connection.ThrowReceiveOnce = true;
            Assert.Throws<TestException>(async () => result = await Client.SendRequestAsync("a"));
        }

        [Test]
        public void EventIsRaisedOnEventPacket()
        {
            int expectedEventDataCount = 2;
            int actualEventDataCount = 0;
            string expectedEventName = "testEvent";
            string actualEventName = null;

            Client.Connect();
            Client.RconEvent += (sender, e) =>
            {
                actualEventDataCount = e.EventData.Count;
                actualEventName = e.EventName;
            };

            Connection.EnqueueEvent(expectedEventName, "foo", "bar");

            SpinWait.SpinUntil(() => actualEventName != null, 1000);

            if (actualEventName == null) Assert.Fail("Timeout");

            Assert.AreEqual(expectedEventName, actualEventName);
            Assert.AreEqual(expectedEventDataCount, actualEventDataCount);
        }

        [Test]
        public void LogsOnCorrectly()
        {
            Client.Connect();

            string pass = "testpass1";
            string salt = "ea702ba4205cb37a88cc84851690a7a5";

            Connection.Password = pass;
            Connection.Salt = salt;

            Client.LogOn(pass);

            Assert.True(Client.IsLoggedOn);
            Assert.True(Connection.LoggedIn);
        }

        [Test]
        public void LogOnThrowsWhenPasswordNotSetOnServer()
        {
            Client.Connect();

            Assert.Throws<PasswordNotSetException>(() => Client.LogOn("dontmatter"));
        }

        [Test]
        public void LogOnThrowsWithInvalidPassword()
        {
            Client.Connect();

            string pass = "testpass1";
            string salt = "ea702ba4205cb37a88cc84851690a7a5";

            Connection.Password = pass;
            Connection.Salt = salt;

            Assert.Throws<IncorrectPasswordException>(() => Client.LogOn("invalid"));
        }

        [Test]
        public void SendingDuringEventDoesNotDeadlock()
        {
            bool eventDone = false;
            bool timedOut = false;

            Client.Connect();
            Client.RconEvent += (sender, e) =>
            {
                IList<string> result;

                try
                {
                    result = Client.SendRequest(1000, "test");
                }
                catch (TimeoutException)
                {
                    timedOut = true;
                }

                eventDone = true;
            };

            Connection.EnqueueEvent("event", "param");

            SpinWait.SpinUntil(() => eventDone, 2000);

            Assert.True(eventDone, "Spin timed out");
            Assert.False(timedOut, "Request timed out");
        }

        private class TestException : Exception { }

        private class TestConnection2 : IRconConnection<Packet>
        {
            private Queue<Packet> PacketQueue = new Queue<Packet>();
            public bool ThrowSend;
            public bool ThrowReceiveOnce;
            public bool LastIsFromClient;
            public bool LastIsResponse;
            public bool AttemptedConnect;
            public bool AttemptedDisconnect;
            public bool LoggedIn;
            public string Password = "";
            public string Salt = "";
            public int Delay = 0;

            public void SendPacket(Packet packet)
            {
                CheckHeader(packet);

                if (ThrowSend)
                    throw new TestException();

                if (packet.Words[0].Equals("login.hashed"))
                {
                    if (string.IsNullOrEmpty(Password))
                    {
                        EnqueueResponse(packet, "PasswordNotSet");
                    }
                    else if (packet.Words.Count == 1)
                    {

                        EnqueueResponse(packet, "OK", Salt);
                    }
                    else if (packet.Words.Count == 2)
                    {
                        if (packet.Words[1].Equals(GeneratePasswordHash(Password, Salt)))
                        {
                            LoggedIn = true;
                            EnqueueResponse(packet, "OK");
                        }
                        else
                        {
                            EnqueueResponse(packet, "InvalidPasswordHash");
                        }

                    }
                    return;
                }

                EnqueueResponse(packet);
            }

            //from bf3 lib
            private static string GeneratePasswordHash(string pass, string salt)
            {
                using (System.Security.Cryptography.MD5 md5 = System.Security.Cryptography.MD5.Create())
                {
                    byte[] passBytes = Encoding.Default.GetBytes(pass);

                    int saltByteCount = salt.Length / 2; //hex string
                    byte[] allBytes = new byte[passBytes.Length + saltByteCount];

                    for (int i = 0; i < salt.Length; i += 2)
                        allBytes[i / 2] = Convert.ToByte(salt.Substring(i, 2), 16);//converts the hex value into a byte

                    for (int i = 0; i < passBytes.Length; i++)
                        allBytes[i + saltByteCount] = passBytes[i];

                    byte[] resultBytes = md5.ComputeHash(allBytes);
                    StringBuilder sb = new StringBuilder();
                    foreach (byte b in resultBytes)
                    {
                        sb.AppendFormat("{0:X2}", b);
                    }

                    return sb.ToString();
                }
            }

            public async Task SendPacketAsync(Packet packet)
            {
                SendPacket(packet);
            }

            public Packet ReceivePacket()
            {
                throw new NotImplementedException();
            }

            public async Task<Packet> ReceivePacketAsync()
            {
                await TaskEx.Delay(Delay);

                if (ThrowReceiveOnce)
                {
                    ThrowReceiveOnce = false;
                    throw new TestException();
                }

                while (PacketQueue.Count == 0)
                    await TaskEx.Delay(200);

                return PacketQueue.Dequeue();
            }

            public virtual void Connect()
            {
                AttemptedConnect = true;
                IsConnected = true;
            }

            public virtual void Disconnect()
            {
                AttemptedDisconnect = true;
                IsConnected = false;
            }

            public virtual bool IsConnected
            {
                get;
                set;
            }

            private void CheckHeader(Packet packet)
            {
                LastIsFromClient = packet.IsFromClient;
                LastIsResponse = packet.IsResponse;
            }

            private void EnqueueResponse(Packet request, params string[] words)
            {
                PacketQueue.Enqueue(new Packet(false, true, request.Sequence, words.Length == 0 ? new[] { "OK" } : words));
            }

            private int EventCounter = Packet.MaxSequenceValue;
            public void EnqueueEvent(params string[] words)
            {
                PacketQueue.Enqueue(new Packet(false, false, --EventCounter, words));
            }

            public void ResetFlags()
            {
                ThrowSend = false;
                ThrowReceiveOnce = false;
                LastIsFromClient = false;
                LastIsResponse = false;
                AttemptedConnect = false;
                AttemptedDisconnect = false;
                LoggedIn = false;
            }
        }

        //correct responses have word "correct"
        private class TestConnection : IRconConnection<Packet>
        {
            //these will be send before correct response
            public Queue<string[]> IncorrectDataToSend = new Queue<string[]>();
            private Queue<int> SequencesToReply = new Queue<int>();
            private Queue<Packet> EventPackets = new Queue<Packet>();
            public bool LastIsFromClient;
            public bool LastIsResponse;
            public bool AttemptedConnect;
            public bool AttemptedDisconnect;
            public bool ThrowSend;
            public bool ThrowReceiveOnce;
            public int Delay = 0;
            public string LoginSalt = string.Empty;
            public string LoginPass = string.Empty;
            private int SendSaltSequence;
            public bool LoggedIn;
            private int SendLoggedInSequence;

            public TestConnection()
            {
                ResetFlags();
            }

            public virtual void SendPacket(Packet packet)
            {
                CheckHeader(packet);

                if (ThrowSend) throw new TestException();

                if (packet.Words[0].Equals("login.hashed"))
                {
                    if (packet.Words.Count == 1)
                    {
                        SendSaltSequence = packet.Sequence;
                    }
                    else if (packet.Words.Count == 2)
                    {
                        if (packet.Words[1].Equals(GeneratePasswordHash(LoginPass, LoginSalt), StringComparison.OrdinalIgnoreCase))
                        {
                            LoggedIn = true;
                            SendLoggedInSequence = packet.Sequence;
                        }
                    }
                    else
                        Assert.Fail("Bad login process.");

                    return;
                }

                SequencesToReply.Enqueue(packet.Sequence);
            }

            //from bf3 lib
            private static string GeneratePasswordHash(string pass, string salt)
            {
                using (System.Security.Cryptography.MD5 md5 = System.Security.Cryptography.MD5.Create())
                {
                    byte[] passBytes = Encoding.Default.GetBytes(pass);

                    int saltByteCount = salt.Length / 2; //hex string
                    byte[] allBytes = new byte[passBytes.Length + saltByteCount];

                    for (int i = 0; i < salt.Length; i += 2)
                        allBytes[i / 2] = Convert.ToByte(salt.Substring(i, 2), 16);//converts the hex value into a byte

                    for (int i = 0; i < passBytes.Length; i++)
                        allBytes[i + saltByteCount] = passBytes[i];

                    byte[] resultBytes = md5.ComputeHash(allBytes);
                    StringBuilder sb = new StringBuilder();
                    foreach (byte b in resultBytes)
                    {
                        sb.AppendFormat("{0:x2}", b);
                    }

                    return sb.ToString();
                }
            }

            public virtual async Task SendPacketAsync(Packet packet)
            {
                CheckHeader(packet);

                if (ThrowSend) throw new TestException();

                SequencesToReply.Enqueue(packet.Sequence);

                return;
            }

            //prob wont be used by this rconclient
            public virtual Packet ReceivePacket()
            {
                throw new NotImplementedException();
            }

            public async virtual Task<Packet> ReceivePacketAsync()
            {
                int t;

                if (SendLoggedInSequence > -1)
                {
                    t = SendLoggedInSequence;
                    SendLoggedInSequence = -1;
                    return new Packet(false, true, t, "OK", LoginSalt);
                }

                if (SendSaltSequence > -1)
                {
                    t = SendSaltSequence;
                    SendSaltSequence = -1;
                    return new Packet(false, true, t, "OK", LoginSalt);
                }

                while (SequencesToReply.Count == 0 && EventPackets.Count == 0)
                {
                    await Task.Delay(500);
                }
                //assume events for incorrect data
                if (IncorrectDataToSend.Count > 0) return new Packet(false, false, 1337, IncorrectDataToSend.Dequeue());
                else if (SequencesToReply.Count > 0)
                {
                    if (Delay > 0)
                        await Task.Delay(Delay);


                    if (ThrowReceiveOnce)
                    {
                        ThrowReceiveOnce = false;
                        throw new TestException();
                    }

                    return new Packet(false, true, SequencesToReply.Dequeue(), "correct");
                }
                else //event
                {
                    return EventPackets.Dequeue();
                }
            }

            public virtual void Connect()
            {
                AttemptedConnect = true;
                IsConnected = true;
            }

            public virtual void Disconnect()
            {
                AttemptedDisconnect = true;
                IsConnected = false;
            }

            public virtual bool IsConnected
            {
                get;
                set;
            }

            private void CheckHeader(Packet packet)
            {
                LastIsFromClient = packet.IsFromClient;
                LastIsResponse = packet.IsResponse;
            }

            public void ResetFlags()
            {
                LastIsFromClient = false;
                LastIsResponse = false;
                AttemptedConnect = false;
                AttemptedDisconnect = false;
                ThrowSend = false;
                ThrowReceiveOnce = false;
                SendSaltSequence = -1;
                LoggedIn = false;
                SendLoggedInSequence = -1;
            }

            public void SendEvent(params string[] words)
            {
                Packet packet = new Packet(false, false, 1, words);

                EventPackets.Enqueue(packet);
            }
        }
    }
}
